#pragma once
#include "auxiliary.h"
#include "Joint.h"
#include "Force.h"
#include "json.h"
class Equation;
class Solver;
class MBSystem
{
private:
	Equation* peq = nullptr;
	Solver* psolver = nullptr;
	unsigned int dof;
	unsigned int nb;
	unsigned int nc;
	vector<JointBase*> jointvec;
	vector<ForceBase*> forcevec;
	vector<JointForce*> jforcevec;
	vector<vector<MatR3CX>> alpha;
	vector<vector<MatR3CX>> beta;
	vector<double> massvec;
	vector<Matrix3d> Jcvec;
	vector<Vector3d> Fo;
	vector<Vector3d> Mo;
	VectorXd Ft;
	vector<Vector3d> w;
	vector<Vector3d> sigma;
	vector<Vector3d> epsilon;
	vector<vector<Vector3d>> d;                             //constant 
	vector<vector<Vector3d>> dstar;
	vector<MatR3CX> Ho;
	vector<MatR3CX> Hh;
	vector<int> L;
	MatrixXi S;
	MatrixXi T;
	MatrixXd Z;
	VectorXd z;
	VectorXd ypre;
	VectorXd dypre;
	vector<int> ydim;
	vector<int> ycum;
	VectorXd q;
	VectorXd dq;
	VectorXd y0;
	VectorXd dy0;
private:
	/*call only in initialization*/

	bool CalMass();
	bool Cald();
	/*call every update*/
	bool CalHo();
	bool CalHh();
	bool CalJc();
	bool Caldstar();
	bool Calepsilon();
	bool Calsigma(IN double t);
	bool CalMo(IN double t);
	bool CalFo(IN double t);
	bool CalFt(IN double t);
	bool Calw(IN double t);
	bool Calbeta();
	bool Calalpha();
	
	
protected:
	bool Initialize();                                                       //call before calculating.
	bool Update(IN double t,IN double* y,IN double* dy);                     //call before func CalZ and Calz.
	bool Calytoq(IN double t,IN double* y);
	bool Caldytodq(IN double t,IN double* dy);
public:
	friend class Equation;
	friend class Solver;
	friend class MBFileParser;
	MBSystem();
	~MBSystem();
	MatrixXd& CalZ(IN double t, IN double* y, IN double* dy);
	VectorXd& Calz(IN double t, IN double* y, IN double* dy);
	bool Add(IN JointBase* j);
	bool Del(IN JointBase* j);
	bool Add(IN ForceBase* f);
	bool Del(IN ForceBase* f);
	bool Add(IN JointForce* jf);
	bool Del(IN JointForce* jf);
	bool Sety0(IN VectorXd& _y0);
	bool Setdy0(IN VectorXd& _dy0);
	bool SetTimeInterval(IN double ti, IN double te, IN int N);
	bool setTolerance(IN double r = 1e-4, IN double a = 1e-3);
	bool SaveAs(IN string fname, IN bool isbinary = false);
	unsigned int DOF()const;
	bool Calculate();
	bool SetToInitial();
};

class Equation
{
	MBSystem* pmbs = nullptr;
	unsigned int dof;                                                        //dimension of the equation
protected:
	MatrixXd L;
	VectorXd R;
public:
	friend class Solver;
	Equation(IN MBSystem* p);
	~Equation();
	bool Initialize();
	MatrixXd& Left(IN double t, IN VectorXd& y);
	VectorXd& Right(IN double t, IN VectorXd& y);
	VectorXd InitialValue()const;
	unsigned int DOF()const;
};

class Solver
{
	Equation* peq = nullptr;
protected:
	double t_ini;
	double t_end;
	unsigned int Nstep;
	double Atol;
	double Rtol;
	vector<double> tspan;
	list<VectorXd> Y;
	list<VectorXd> DY;
	list<VectorXd> Q;
	list<VectorXd> DQ;
public:
	friend class MBSystem;
	friend class MBFileParser;
	Solver(IN Equation* p);
	~Solver();
	bool SetTime(IN double ti,IN double te,IN unsigned int N);
	bool SetTolerance(IN double atol,IN double rtol);
	bool Calculate();
};

class MBFileParser
{
	MBSystem* pmbs = nullptr;
	vector<Body*> bodyvec;
	bool freememo = false;
protected:
	void clear();
	void CheckId(IN int id);
	void CheckMass(IN double m);
	void CheckJc(IN const Json::Value& Jc);
	void CheckRho(IN const Json::Value& rho);
	void CheckMat3d(IN const Json::Value& val);
	void CheckPos(IN const Json::Value& val, IN unsigned int k);
	void CheckVel(IN const Json::Value& val, IN unsigned int k);
	void CheckForce(IN const Json::Value& val, IN unsigned int k);
	void GetJc(IN const Json::Value& Jc, OUT Matrix3d& Ic);
	void GetRho(IN const Json::Value& rho, OUT Vector3d& r);
	void GetMat3d(IN const Json::Value& val, OUT Matrix3d& M);
	void GetPos(IN const Json::Value& val, OUT VectorXd& p);
	void GetVel(IN const Json::Value& val, OUT VectorXd& v);
	void GetForce(IN const Json::Value& val, OUT VectorXd& f);
public:
	MBFileParser() :pmbs(nullptr) {};
	~MBFileParser();
	bool Read(IN const string& fname);
	bool Write(IN const string& fname);
	bool Simulate();
	bool SaveDataAs(IN const string& fname, IN bool isbinary);
};